#include "gst/gstbin.h"
#include "gst/gstbus.h"
#include "gst/gstclock.h"
#include "gst/gstelement.h"
#include "gst/gstmessage.h"
#include "gst/gstobject.h"
#include "gst/gstpad.h"
#include "gst/gstparse.h"
#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <gst/gst.h>

#include "classid.h"
#include "meta_utils.h"

typedef struct Context {
    guint count;
    guint frame_count;
    GstElement *textoverlay;
    ClassIdIndex *classid;
} Context;

static gboolean test_timeout(gpointer arg) {
    Context *ctx = (Context *)arg; 
    char string_buffer[128];
    g_printerr("time out fired %u\n", ctx->count);
    sprintf(string_buffer, "count %d", ctx->count);
    ctx->count += 1;
    g_object_set(G_OBJECT(ctx->textoverlay), "text", string_buffer, NULL);
    return TRUE;
}

static void msg_handler(GstBus *bus, GstMessage *msg, GMainLoop *main_loop) {
    if (GST_MESSAGE_TYPE(msg) == GST_MESSAGE_EOS) {
        printf("stream normal end\n");
        g_main_loop_quit(main_loop);
        return;
    }

    if (GST_MESSAGE_TYPE(msg) == GST_MESSAGE_ERROR) {
        printf("stream in a error\n");
        g_main_loop_quit(main_loop);
        return;
    }
}

GstPadProbeReturn simple_probe(GstPad *pad, GstPadProbeInfo *info, gpointer arg) {
    char line_buffer[128];
    Context *ctx = (Context *)arg;
    g_printerr("simple probe called\n");

    if (info->data == NULL) {
        g_printerr("no buffer found\n");
        return GST_PAD_PROBE_OK;
    }

    if (GST_IS_BUFFER(info->data)) {
        GstBuffer *buffer = GST_PAD_PROBE_INFO_BUFFER(info);
        g_printerr("probe found a buffer #%d:\n", ctx->frame_count++);
        // print_buffer_meta(buffer);
        // const char *text = get_buffer_class(buffer);
        const NvDsLabelInfo *info = get_buffer_label(buffer);
        if (info != NULL) {
            sprintf(line_buffer, "%.2f, #%4d, %s", info->result_prob, info->result_class_id, info->result_label);
            g_printerr("infer class: %s\n", line_buffer);
            g_object_set(G_OBJECT(ctx->textoverlay), "text", line_buffer, NULL);

            TensorView tensor = get_buffer_tensor(buffer);
            if (tensor.data != NULL) {
                g_print("tensor fetched, tensor dim:(");
                for (int i = 0; i < tensor.dims->numDims; ++i) {
                    g_print("%d,", tensor.dims->d[i]);
                } g_print(")\n");
            }
            const float *arr = tensor.data;
            float tensor_prob = arr[info->result_class_id];
            g_print("tensor_prob: %.5f, result_prob: %.5f\n", tensor_prob, info->result_prob);
        } else {
            g_object_set(G_OBJECT(ctx->textoverlay), "text", "", NULL);
        }

    } else {
        GstBufferList *buffer_list = GST_PAD_PROBE_INFO_BUFFER_LIST(info);
        guint len = gst_buffer_list_length(buffer_list);
        const char *text = NULL;
        for (guint i = 0; i < len; ++i) {
            GstBuffer *buffer = gst_buffer_list_get(buffer_list, i);

            if (buffer == NULL) continue;
            g_printerr("probe found a buffer in list[%d], #%d:", i, ctx->frame_count++);
            text = get_buffer_class(buffer);
            if (text != NULL) {
                g_printerr("infer class: %s\n", text);
            }
        }
        if (text != NULL) {
            g_object_set(G_OBJECT(ctx->textoverlay), "text", text, NULL);
        } else {
            g_object_set(G_OBJECT(ctx->textoverlay), "text", "", NULL);
        }
    }

    return GST_PAD_PROBE_OK;
}

int gst_main(int argc, char **argv, ClassIdIndex *classid) {

    Context ctx;
    ctx.count = 0;
    ctx.frame_count = 1;
    ctx.classid = classid;
    ctx.textoverlay = NULL;

    GstElement *pipeline;
    GstBus *bus;
    GMainLoop *timeout_loop = g_main_loop_new(NULL, FALSE);
    gulong probe_id = 0;

    gst_init(&argc, &argv);

    pipeline = gst_parse_launch(
        "v4l2src ! image/jpeg, width=800, height=600, framerate=30/1, format=(string)MJPG "
        "! jpegdec "
        "! video/x-raw, format=(string)I420 "
        "! tee name=v_src "
        "! nvvidconv flip-method=4 "
        "! video/x-raw(memory:NVMM), format=(string)NV12 "
        "! mux.sink_0 nvstreammux name=mux batch-size=1 width=224 height=224 live-source=1 sync-inputs=0 "
        "! nvinfer unique-id=1 config-file-path=/home/baomingkun/Workspace/gstreamer-c-example/infer-config.txt "
        "! fakesink name=infer-sink "
        "v_src.  ! videoflip method=horizontal-flip "
        "! textoverlay name=subtitle "
        "! xvimagesink",
        NULL
    );

    GstElement *textoverlay = gst_bin_get_by_name(GST_BIN(pipeline), "subtitle");
    ctx.textoverlay = textoverlay;

    GstElement *infer_sink = gst_bin_get_by_name(GST_BIN(pipeline), "infer-sink");

    GstPad *infer_sink_sink = gst_element_get_static_pad(infer_sink, "sink");
    if (infer_sink_sink == NULL) {
        g_printerr("can't get overlay sink pad\n");
        // gst_object_unref(textoverlay);
        gst_object_unref(pipeline);
        gst_object_unref(infer_sink);
        g_main_loop_unref(timeout_loop);
        return -1;
    }
    probe_id = gst_pad_add_probe(
        infer_sink_sink, 
        GST_PAD_PROBE_TYPE_BUFFER | GST_PAD_PROBE_TYPE_BUFFER_LIST | GST_PAD_PROBE_TYPE_PULL | GST_PAD_PROBE_TYPE_PUSH,
        simple_probe, &ctx, NULL
    );


    bus = gst_element_get_bus(pipeline);

    gst_bus_add_signal_watch(bus);
    g_signal_connect(bus, "message", G_CALLBACK(msg_handler), timeout_loop);

    gst_element_set_state(pipeline, GST_STATE_PLAYING);

    // set main loop to run
    g_main_loop_run(timeout_loop);
    g_main_loop_unref(timeout_loop);

    // main loop run finish
    g_print("finish main loop, cleanup\n");

    gst_pad_remove_probe(infer_sink_sink, probe_id);
    gst_object_unref(infer_sink_sink);
    gst_object_unref(infer_sink);

    // gst_message_unref(msg);
    gst_object_unref(bus);
    gst_element_set_state(pipeline, GST_STATE_NULL);
    gst_object_unref(pipeline);

    return 0;
}

int main(int argc, char **argv) {

    printf("hello world\n");


    const char *classid_path = "./index.json";
    if (argc > 1) {
        classid_path = argv[1];
    }
    ClassIdIndex *classid = readClassIdIndex(classid_path);
    if (classid == NULL) {
        return -1;
    }

    printf("index[0]: %s\n", getClassName(classid, 0));
    printf("index[1]: %s\n", getClassName(classid, 1));

    assert(getClassName(classid, 2) == NULL);

    gst_main(argc, argv, classid);

    freeClassIdIndex(classid);

    return 0;
}